/**
 * Simple cmd_vel demo for unitree go2 robot
 * Subscribe to /cmd_vel and send move commands to robot
 **/

use r2r::{self, QosProfile};
use serde_json::json;
use std::sync::atomic::{AtomicBool, Ordering};
use futures::{executor::LocalPool, stream::StreamExt, task::LocalSpawnExt};

static RUNNING: AtomicBool = AtomicBool::new(true);

// Safety limits
const MAX_LINEAR_X: f32 = 0.8;    // m/s
const MAX_LINEAR_Y: f32 = 0.4;    // m/s
const MAX_ANGULAR_VEL: f32 = 1.0; // rad/s

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Setup Ctrl+C handler
    ctrlc::set_handler(move || {
        println!("\nReceived Ctrl+C, shutting down...");
        RUNNING.store(false, Ordering::SeqCst);
    })?;

    println!("Starting simple cmd_vel demo...");
    
    // Create ROS2 context and node
    let ctx = r2r::Context::create()?;
    let mut node = r2r::Node::create(ctx, "simple_cmd_vel", "")?;
    
    // Create publisher for sport requests
    let req_puber = node.create_publisher::<r2r::unitree_api::msg::Request>(
        "/api/sport/request",
        QosProfile::default(),
    )?;
    
    // Create subscriber for cmd_vel
    let cmd_vel_subscriber = node.subscribe::<r2r::geometry_msgs::msg::Twist>(
        "/cmd_vel",
        QosProfile::default(),
    )?;
    
    println!("✓ ROS2 setup complete. Listening for /cmd_vel messages...");
    println!("Press Ctrl+C to exit");
    
    // Setup async task pool
    let mut pool = LocalPool::new();
    let spawner = pool.spawner();
    
    // Clone publisher for async task
    let req_puber_clone = req_puber.clone();
    
    // Spawn async task to handle cmd_vel messages
    spawner.spawn_local(async move {
        let mut sub = cmd_vel_subscriber;
        while let Some(msg) = sub.next().await {
            if !RUNNING.load(Ordering::SeqCst) {
                break;
            }
            handle_cmd_vel(&msg, &req_puber_clone);
        }
    })?;
    
    // Main loop
    while RUNNING.load(Ordering::SeqCst) {
        node.spin_once(std::time::Duration::from_millis(10));
        pool.run_until_stalled();
        std::thread::sleep(std::time::Duration::from_millis(10));
    }
    
    // Send stop command before exit
    send_stop_command(&req_puber);
    println!("Demo stopped.");
    
    Ok(())
}

fn handle_cmd_vel(twist_msg: &r2r::geometry_msgs::msg::Twist, req_puber: &r2r::Publisher<r2r::unitree_api::msg::Request>) {
    // Extract and limit velocities
    let vx = (twist_msg.linear.x as f32).clamp(-MAX_LINEAR_X, MAX_LINEAR_X);
    let vy = (twist_msg.linear.y as f32).clamp(-MAX_LINEAR_Y, MAX_LINEAR_Y);
    let vyaw = (twist_msg.angular.z as f32).clamp(-MAX_ANGULAR_VEL, MAX_ANGULAR_VEL);
    
    // Log command
    if vx != 0.0 || vy != 0.0 || vyaw != 0.0 {
        println!("Move: vx={:.2}, vy={:.2}, vyaw={:.2}", vx, vy, vyaw);
    }
    
    // Send move command
    send_move_command(vx, vy, vyaw, req_puber);
}

fn send_move_command(vx: f32, vy: f32, vyaw: f32, req_puber: &r2r::Publisher<r2r::unitree_api::msg::Request>) {
    let mut req = r2r::unitree_api::msg::Request::default();
    
    // Create move command
    let js = json!({"x": vx, "y": vy, "z": vyaw});
    req.parameter = js.to_string();
    req.header.identity.api_id = 1008; // MOVE command
    
    // Publish
    if let Err(e) = req_puber.publish(&req) {
        eprintln!("Failed to send move command: {:?}", e);
    }
}

fn send_stop_command(req_puber: &r2r::Publisher<r2r::unitree_api::msg::Request>) {
    println!("Sending stop command...");
    send_move_command(0.0, 0.0, 0.0, req_puber);
}